package com.xiva.zk.curator;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;

/**
 * @author xiva
 * @Description: zookeeper 分布式独占锁
 */

public class ZkDistributedLock implements Lock, Watcher
{
    private static final Logger logger = LoggerFactory.getLogger(ZkDistributedLock.class);
    private ZooKeeper zk;
    private String root = "/locks";//根
    private String lockName;//竞争资源的标志

    private String waitNode;//等待前一个锁

    private String myZnode;//当前锁

    private volatile CountDownLatch latch;//计数器

    private List<Exception> exception = new ArrayList<Exception>();

    private final static String splitStr = "_lock_";

    private final static long defaultWaitTime = 30 * 1000;

    /**
     * 创建分布式锁,使用前请确认config配置的zookeeper服务可用
     *
     * @param config   127.0.0.1:2181
     * @param lockName 竞争资源标志,lockName中不能包含单词lock
     */
    public ZkDistributedLock(String config, String lockName)
    {
        this.lockName = lockName;
        // 创建一个与服务器的连接
        try
        {
            zk = new ZooKeeper(config, 60 * 1000, this);
            Stat stat = zk.exists(root, false);
            if (stat == null)
            {
                // 创建根节点
                zk.create(root, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
        }
        catch (IOException e)
        {
            exception.add(e);
        }
        catch (KeeperException e)
        {
            exception.add(e);
        }
        catch (InterruptedException e)
        {
            exception.add(e);
        }
    }

    /**
     * zookeeper节点的监视器
     */
    public void process(WatchedEvent event)
    {
        if (this.latch != null)
        {
            this.latch.countDown();
        }
    }

    /*
     * 申请锁的默认时间是 defaultWaitTime ，超时后抛出LockException运行时异常
     */
    public void lock()
    {
        try
        {
            if (this.tryLock())
            {
                logger.debug("Thread " + Thread.currentThread().getId() + " " + myZnode + " get lock true");
                return;
            }
            else
            {
                waitForLock(waitNode, defaultWaitTime, TimeUnit.MILLISECONDS, true);//等待锁
            }
        }
        catch (KeeperException e)
        {
            throw new LockException(e);
        }
        catch (InterruptedException e)
        {
            throw new LockException(e);
        }
    }

    /*
     * 非阻塞方式获取锁
     * 如果存在锁返回false
     * 通过判断返回值判断是否获取锁成功
     */
    public boolean tryLock()
    {
        if (exception.size() > 0)
        {
            throw new LockException(exception.get(0));
        }
        try
        {
            if (lockName.contains(splitStr)) throw new LockException("lockName can not contains \\u000B");
            //创建临时子节点
            myZnode = zk.create(root + "/" + lockName + splitStr, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            logger.info(myZnode + " is created ");
            List<String> lockObjNodes = this.getChildrenListSort();
            logger.debug(myZnode + "==" + lockObjNodes.get(0));
            if (myZnode.equals(root + "/" + lockObjNodes.get(0)))
            {
                //如果是最小的节点,则表示取得锁
                return true;
            }
            //如果不是最小的节点，找到比自己小1的节点
            String subMyZnode = myZnode.substring(myZnode.lastIndexOf("/") + 1);
            waitNode = lockObjNodes.get(Collections.binarySearch(lockObjNodes, subMyZnode) - 1);
        }
        catch (KeeperException e)
        {
            throw new LockException(e);
        }
        catch (InterruptedException e)
        {
            throw new LockException(e);
        }
        return false;
    }

    /*
     * 非阻塞方式获取锁
     * 通过判断返回值判断是否获取锁成功
     */
    public boolean tryLock(long time, TimeUnit unit)
    {
        try
        {
            if (this.tryLock())
            {
                return true;
            }
            return waitForLock(waitNode, time, unit, false);
        }
        catch (Exception e)
        {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * @param lower    次小节点
     * @param waitTime 等待时间
     * @param unit     时间单位
     * @param isBreak  超时是否抛出异常
     * @return
     * @throws InterruptedException
     * @throws KeeperException
     */
    private boolean waitForLock(String lower, long waitTime, TimeUnit unit, boolean isBreak) throws InterruptedException, KeeperException
    {
        long beginTime = System.currentTimeMillis();
        Stat stat = zk.exists(root + "/" + lower, true);
        //判断比自己小一个数的节点是否存在,如果不存在则无需等待锁,同时注册监听
        if (stat != null)
        {
            logger.debug("Thread " + Thread.currentThread().getId() + " waiting for " + root + "/" + lower);
            this.latch = new CountDownLatch(1);
            this.latch.await(waitTime, unit);
            if (this.latch.getCount() == 1)
            {
                this.latch = null;
                if (isBreak)
                {
                    throw new LockException("申请锁资源默认时间是：" + defaultWaitTime + "ms,等待锁超时");
                }
                else
                {
                    return false;
                }
            }
            this.latch = null;
        }

        //取出所有lockName的锁
        List<String> lockObjNodes = this.getChildrenListSort();
        logger.debug(myZnode + "==" + lockObjNodes.get(0));
        if (myZnode.equals(root + "/" + lockObjNodes.get(0)))
        {
            //如果是最小的节点,则表示取得锁
            return true;
        }
        else
        {
            //如果不是最小的节点，找到比自己小1的节点
            String subMyZnode = myZnode.substring(myZnode.lastIndexOf("/") + 1);
            waitNode = lockObjNodes.get(Collections.binarySearch(lockObjNodes, subMyZnode) - 1);
            long endTime = System.currentTimeMillis();
            waitTime = waitTime - (endTime - beginTime);
            if (waitTime <= 0) return false;
            return waitForLock(lower, waitTime, unit, false);//等待锁
        }
    }

    /**
     * 取出本次锁的所有子节点（升序）
     *
     * @return
     * @throws LockException
     */
    private List<String> getChildrenListSort() throws KeeperException, InterruptedException
    {
        //取出所有子节点
        List<String> subNodes = zk.getChildren(root, false);
        //取出所有lockName的锁
        List<String> lockObjNodes = new ArrayList<String>();
        for (String node : subNodes)
        {
            String _node = node.split(splitStr)[0];
            if (_node.equals(lockName))
            {
                lockObjNodes.add(node);
            }
        }
        Collections.sort(lockObjNodes);
        return lockObjNodes;
    }

    public void unlock()
    {
        try
        {
            logger.debug("unlock " + myZnode);
            if (null != myZnode)
            {
                zk.delete(myZnode, -1);
            }
            myZnode = null;
//            zk.close();
        }
        catch (InterruptedException e)
        {
            e.printStackTrace();
        }
        catch (KeeperException e)
        {
            e.printStackTrace();
        }
    }

    public void releaseLockClient()
    {
        try
        {
            zk.close();
        }
        catch (InterruptedException e)
        {
            e.printStackTrace();
        }
    }

    public void lockInterruptibly() throws InterruptedException
    {
        this.lock();
    }

    public Condition newCondition()
    {
        return null;
    }

    public class LockException extends RuntimeException
    {
        private static final long serialVersionUID = 1L;

        public LockException(String e)
        {
            super(e);
        }

        public LockException(Exception e)
        {
            super(e);
        }
    }
}
